#!/usr/bin/env python

"""
If called from a Unix shell, please provide the function name you want to
call, and any required arguments, exactly as presented here.

e.g.::

    python count_hmms.py countseqs /path/to/dir

or::

    python count_hmms.py

This will traverse all directories from Archaea, Bacteria and Eukaryota,
and will return the total number of hmms that exist and the number that
are in the right place.


Description:

 - Some functions to look through the SSUMMO directory structure,
   checking either the HMMs or the sequences found in that directory.
 - If a function argument is presented here with a default value, then
   that argument is optional. Default values are the ones with an equal
   sign.

e.g. The function 'walkall' has an argument 'delete', which has a
  default value of False.

Defined functions:-
  walkall( dir_list = ['Bacteria','Eukaryota','Archaea'] , delete=False )
  countseqs(pathname, file_name='accessions.txt'):
  avseqlens( pathname ):
  os_walk(dir_list, delete=False)
  counthmms(cwd):

Other help??:-
  pydoc count_hmms.py
"""

#
# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>
#

import os
import sys
import subprocess
from ssummo import CONFIG

def walkall(dir_list = None , delete=False):
    if dir_list is None:
        dir_list = ['Bacteria', 'Eukaryota', 'Archaea']
    """Walks all directories, tallying correctly placed hmm files.
    By default, don't delete misplaced hmms."""
    os_walk(dir_list, delete)

def countseqs( pathname , file_name = 'accessions.txt'):
    """Given a pathname to a taxonomic id,
    Return the number of representative 16S sequences in a
    taxonomies folder directory."""
    pathname = pathname.rstrip()
    if pathname[-1] == '/':
        pathname = pathname[:-1]
    if file_name == 'accessions.txt':
        try:
            handle = file(os.path.join(pathname, 'accessions.txt'), 'r')
        except IOError:
            print("No accessions file at:- {0}"\
                  .format(os.path.join(pathname, 'accessions.txt')))
            return 0
    else:
        with file( file_name ,'r') as sequence_file:
            seq_count = 0
            for line in sequence_file.xreadlines():
                if line.startswith('>'):
                    seq_count += 1
        return seq_count
    seq_count = 1
    for char in handle.readline():
        if char == ' ':
            seq_count += 1
    return seq_count

def avseqlens( pathname ):
    """Given a pathname, find the sequences file, and average the 
    length of all the sequences"""
    from Bio import SeqIO, AlignIO
    tot_seq_lens = 0
    nseqs = 0
    try:
        handle = file( os.path.join(pathname, 'sequences.fa'), 'r')
        for sequence in SeqIO.parse(handle,'fasta'):
            tot_seq_lens += len(sequence)
            nseqs += 1
    except OSError:
        cmd = '{fastacmd} -d {top} -i {_input} -o {output}' \
              .format(fastacmd=os.path.join(CONFIG.blastDir, 'fastacmd' ),
                      top=CONFIG.top,
                      _input=os.path.join(pathname, 'accessions.txt'),
                      output=os.path.join(pathname, 'sequences.fa'))
        subprocess.call(cmd)
        handle = file(os.path.join(pathname, 'sequences.fa'), 'r')
        for sequence in SeqIO.parse(handle, 'fasta'):
            tot_seq_lens += len(sequence)
            nseqs += 1
    handle.close()
    av_seq_len = tot_seq_lens / nseqs
    print("Average sequence length in {0} is: {1}"\
          .format(pathname, av_seq_len))
    try:
        handle = file(os.path.join(pathname, 'alignment.fa'), 'r')
        print("Figuring out the alignment sequence length of a {0}MB file."\
             .format(os.path.getsize(os.path.join(pathname, 'alignment.fa')) / 1000.))
        alignment = AlignIO.read(handle, 'stockholm')
        print("Alignment length: {0}".format(alignment.get_alignment_length()))
    except IOError:
        print("no alignment file here. Not gonna make one either....")
    return av_seq_len

def counthmms( cwd):
    """Given a pathname, this function will return the number of HMMs in that
    directory, and the number of directories in that directory"""
    contents = os.listdir(cwd)
    hmm_count = 0
    dir_count = 0
    for cont in contents:
        if cont.endswith('.hmm'):
            hmm_count += 1
        elif '.' not in cont:
            dir_count += 1
        else:
            continue
    return hmm_count, dir_count

def os_walk( dir_list, delete = False):
    """Uses os.walk() and traverses the directories presented in the list
    'dir_list' and checks to see if hmms are present in the correct place.
    If delete is True, then these hmms will be deleted.
    If delete is False or not specified, then a count of misplaced hmms
    will be printed & returned.
    """
    hmm_count = 0
    dir_count = 0
    tot_hmm_count = 0
    start_dir = CONFIG.arbDBdir
    print('starting from {0}'.format(CONFIG.arbDBdir))
    if type(dir_list) == str and os.path.exists(dir_list):
        dir_list = [dir_list]
    for _dir in dir_list:
        _dir = os.path.join(start_dir, _dir )

        for root, dirs, files in os.walk(_dir):
            dir_count += 1
            path = root
            end_path = path[path.rfind('/')+1:]
            hmm_true = False    ### Initiate variable that tests if .hmm file in correct place
            for fil in files:
#                   hmm_true = False    ### Initiate variable that tests if .hmm file in correct place.
                if fil.endswith('.hmm'):
                    if fil[:fil.rfind('.')] == end_path:
                        hmm_count += 1
                        hmm_true = True
                    else:
                        if delete == True:
                            os.remove( os.path.join(root,fil) )
                            print("removed '{0}'".format(
                                os.path.join(root , fil) ))
                        else:
                            print("would remove '{0}'".format(
                                os.path.join(root, fil)))
                    tot_hmm_count += 1
                    if hmm_true == False:
                        print("{0} doesn't have hmm in correct place"\
                                .format(os.path.join(root, fil)))
            if hmm_true == False:
                print(root)

    print('num of dirs = {0}'.format(dir_count))
    print('num of hmms in correct place = {0}'.format(hmm_count))
    print('total number of HMMs = {0}'.format(tot_hmm_count))


if __name__ == '__main__':
    dir_list = ['Bacteria', 'Eukaryota', 'Archaea']
    try:
        fn = sys.argv[1]
    except IndexError:
        error = '\n\nWrong number of arguments specified. To see help documentation, try:-'
        error += '\n$ pydoc {0}\n'.format( __file__ )
        sys.stderr.write(error)
        exit(1)
    if len(sys.argv) == 4:
        arg = sys.argv[2]
        arg2 = sys.argv[3]
    elif len(sys.argv) == 3:
        arg = sys.argv[2]
    if fn == 'walkall':
        try:
            walkall(dir_list = arg, delete=False)
        except NameError:
            try:
                walkall(arg)
            except NameError:
                walkall()
    elif fn == 'countseqs':
        print(countseqs(arg))
    elif fn == 'counthmms':
        counthmms(arg)
    elif fn == 'os_walk':
        if len(sys.argv) >= 4:
            os_walk(arg, arg2)
        else:
            os_walk(arg)
    elif fn == 'avseqlens':
        avseqlens(arg)
    else:
        print("Unknown function. Functions are:-")
        print("walkall, countseqs, counthmms & os_walk")
